import sqlite3
from sqlite3 import Error
from flask import current_app as app


def create_files_table():
    try:
        conn = sqlite3.connect('rag.db')
        cursor = conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS files (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                upload_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                file_name TEXT NOT NULL,
                file_type TEXT NOT NULL,
                kb_id INTEGER NOT NULL,
                doc_count INTEGER NOT NULL
            )
        ''')
        conn.commit()
        conn.close()
    except Error as e:
        app.logger.error(f'创建files表出错: {str(e)}')


# 定义知识库表的创建SQL
CREATE_KNOWLEDGE_BASE_TABLE_SQL = '''
CREATE TABLE IF NOT EXISTS knowledge_bases (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    title TEXT NOT NULL,
    remark TEXT,
    sort INTEGER NOT NULL,
    enabled BOOLEAN NOT NULL
);
'''

# 连接到SQLite数据库并创建知识库表

def create_knowledge_base_table():
    conn = sqlite3.connect('rag.db')
    cursor = conn.cursor()
    cursor.execute(CREATE_KNOWLEDGE_BASE_TABLE_SQL)
    conn.commit()
    conn.close()


from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

# 定义合并小段落的函数
def merge_small_paragraphs(paragraphs, model, similarity_threshold=0.7):
    """
    基于语义相似度合并小段落。

    Args:
        paragraphs (list): 包含小段落的列表。
        model (SentenceTransformer): SentenceTransformer 模型实例。
        similarity_threshold (float): 合并的相似度阈值。

    Returns:
        list: 合并后的段落列表。
    """
    embeddings = model.encode(paragraphs)
    merged_paragraphs = []
    i = 0
    while i < len(paragraphs):
        current_paragraph = paragraphs[i]
        current_embedding = embeddings[i]
        j = i + 1
        while j < len(paragraphs):
            next_paragraph = paragraphs[j]
            next_embedding = embeddings[j]
            similarity = cosine_similarity([current_embedding], [next_embedding])[0][0]
            if similarity >= similarity_threshold:
                current_paragraph += " " + next_paragraph  # 将下一个段落合并到当前段落
                current_embedding = model.encode([current_paragraph])[0]  # 重新计算合并后段落的嵌入
                j += 1
            else:
                break  # 如果相似度低于阈值，停止与后续段落合并
        merged_paragraphs.append(current_paragraph)
        i = j
    return merged_paragraphs